/* vim: ft=c */
#ifndef OPALLOC_INCLUDED
#define OPALLOC_INCLUDED
/******************************************************************************
* (c)2022-2025 Michael T. Richter
*
* This software is distributed under the terms of WTFPLv2.  The full terms and
* text of the license can be found at http://www.wtfpl.net/txt/copying
******************************************************************************/
/** @file
 *  @brief An "object pool" framework for minimally-fragmenting interlaced
 *         dynamic allocation and deallocation of dynamic memory objects in C.
 *
 * In embedded and small-memory systems, excessive use of dynamically
 * allocated objects can cause heap fragmentation such that, despite there
 * being sufficient free space in the heap for an object, calls to `malloc()`
 * and related functions fail because no single block is large enough to
 * fulfill the desired request.
 *
 * This library provides an implementation of a smarter memory pool manager that
 * combats this problem by minimizing the number of interlaced allocations,
 * deallocations and reallocations at a slight expense in memory and CPU
 * overhead.  It supplies a low level interface for building custom allocators
 * in cases where fine control is needed, a high level interface for the most
 * common and simple use cases, and a macro interface for the easiest use case.
 *
 * Important philosophical note: this library is built on a "don't supply what
 * you probably don't need" basis.  It is also distributed under WTFPLv2 which
 * means you can do to it whatever you like and I don't care.  The specifics of
 * what I don't care about here include:
 *
 *  1. Pre-emption.  I only do work with main loops, reactors, and coroutines.
 *     I don't do pre-emptive multi-tasking because it's dangerous.  As such I
 *     have not littered my code with hooks for mutexes, critical sections, spin
 *     locks and other such pre-emptive shitfuckery.  If you need these, use the
 *     power of "What The Fuck" from WTFPLv2 and put 'em in.  It's your code now
 *     do with it what you will.
 *  2. Compiler magic.  There's no `restrict` or `__attribute__((malloc))` or
 *     other such magical incantations here.  My target environment is embedded
 *     systems and, let's be honest, these incantations are not going to make any
 *     kind of noticeable difference in a sane embedded environment.  Certainly
 *     not sufficiently noticeable for me to want to start that dangerous dark
 *     path of macromania where I put OP_ATTRIBUTE_MALLOC in front of half the
 *     lines of code that will expand to nothing in most cases anyway.  If you
 *     think these compiler magic details are important, use the power of "What
 *     The Fuck" and add them.
 *  3. Unnecessary optimizations.  `op_ll_get_allocator_stats()`, for example is
 *     O(n).  I could make it O(1) by tracking some things, but that tracking is
 *     excess code that needs testing, can be fragile, and **is not necessary**
 *     because that function exists only to help debug allocation problems in
 *     my code, not to sit inside a hot loop for a billion repetitions per
 *     second.  So I'm not going to do it.  If you want it for some reason, you
 *     know the solution by now.
 *  4. Whining about "professionalism".  The project is licensed literally under
 *     the "What The Fuck Public Licence v2".  I think professionalism has long
 *     been thrown out the window from the 52nd storey.  You're not paying me
 *     for this, so by definition this is not a professional product.  If you
 *     want "professionalism" in my documentation and comments, hire me.  Or
 *     don't.  I'm easy either way.
 *
 *  The ultimate answer to any problem you might have is:
 *
 *                      THE POWER OF WTFPLV2 ALLOWS YOU!
 *
 * @see https://personaljournal.ca/qqmrichter/dynamically-static-allocation-in-embedded-systems
 */

#include <stdbool.h>
#include <stddef.h>

/** @defgroup llinterface Low-level interface
 *
 * This is the nuts-and-bolts interface to the library.  It is primarily used to
 * implement the high level interface and the macro interface but has been
 * exposed to enable end-users to make task-specific interfaces if they so
 * choose.
 *
 * It is strongly recommended not to use this interface directly as it is highly
 * error-prone.
 *
 * @{
 */

#ifndef UNUSED
#define UNUSED(X) ((void)(X))
#endif

/** @brief Allocator mode handling. */
typedef enum op_ll_allocator_mode
{
    OP_DOUBLING_INDIVIDUAL, /**< doubling growth, individual object allocation */
    OP_DOUBLING_CHUNK,      /**< doubling growth, chunk object allocation      */
    OP_LINEAR_INDIVIDUAL,   /**< linear growth, individual object allocation   */
    OP_LINEAR_CHUNK,        /**< linear growth, chunk object allocation        */
} op_ll_allocator_mode;

/** @brief Opaque allocator handle permitting multiple object pools. */
typedef struct _op_allocator *op_allocator;

/** @brief Stats for measuring and debugging allocators. */
typedef struct op_allocator_stats
{
    size_t object_size;     /**< size in bytes of stored objects                   */
    size_t maximum_objects; /**< maximum number of objects currently allocated for */
    size_t active_objects;  /**< number of objects actively in use                 */
} op_allocator_stats;

/** @brief Error callback for allocator errors.
 *
 * @param [in] file          The source file where the error was discovered.
 * @param [in] line          The line of the file where the error was discovered.
 * @param [in] error_message The error's textual message.
 *
 * A function pointer
 */
typedef void (*op_error_handler_t)(const char *file, const int line, const char *error_message);
extern op_error_handler_t op_error_handler;

/** @brief Allocate an object in the given allocator context.
 *
 * @param [in,out] allocator The allocator from which the memory is to be allocated.
 *
 * @return A pointer to the memory space, NULL on failure.
 *
 * @note The `op_error_handler()` callback is called before NULL is returned.
 */
void *op_ll_allocate_object(op_allocator allocator);

/** @brief Deallocate a specified object from the given allocator.
 *
 * @param [in, out] allocator The allocator holding the object being freed.
 * @param [in]      object    The object to be freed.
 */
void op_ll_deallocate_object(op_allocator allocator, const void *object);

/** @brief De-initialize an allocator, freeing any owned resources.
 *
 * @param [in, out] allocator The allocator from which to free the object.
 *
 * @note After returning, the allocator object is invalid and can no longer be
 *       used.
 */
void op_ll_deinitialize_allocator(op_allocator allocator);

/** @brief Collect some allocator stats for measurement and debugging.
 *
 * @param [in] allocator The allocator to collect stats from.
 *
 * @return A structure with some allocator statistics for measurement and
 *         debugging.
 */
op_allocator_stats op_ll_get_allocator_stats(const op_allocator allocator);

/** @brief Initialize an allocator according to the provided configuration.
 *
 * @param [in] object_size   The size each individual object requires in RAM.
 * @param [in] initial_count The starting size of the allocator's array.
 * @param [in] use_chunks    Signals this allocator uses chunk allocation.
 * @param [in] use_linear    Signals this allocator uses linear growth.
 *
 * @return An `op_allocator` handle used in subsequent operations or NULL on
 *         failure.
 *
 * @note 1. If `use_chunks` is not provided, only the indirect array is
 *          calloced in the allocator.  If `use_chunks` is provided, the first
 *          `initial_count` objects of `object_size` bytes are also allocated.
 *       2. On failure NULL is returned, but `op_error_handler()` is called
 *          first.
 */
op_allocator op_ll_initialize_allocator(const size_t object_size, const size_t initial_count,
                                        const op_ll_allocator_mode mode);

/**@}*/

/** @defgroup hlinterface High-level (Macro) interface
 *
 * @param [in] TYPE  The type of objects to be allocated.
 * @param [in] COUNT The initial object count for the allocator.
 * @param [in] MODE  The mode of the allocator (as per the low-level interface).
 *
 * The high-level interface is the intended use of this library.  The focus is
 * on a component based design where a single source file contains a single
 * component which is allocated and deallocated dynamically, thus profiting from
 * reuse to minimize fragmentation and increase amortized allocation speed.
 *
 * Each use of the `OP_HL_DECLARE_ALLOCATOR()` macro creates one type-safe
 * allocator using the starting count and mode provided.  It also creates
 * several static functions whose name invoke the given type.  For example:
 *
 *     OP_HL_DECLARE_ALLOCATOR(my_fancy_type, 4, OP_DOUBLING_CHUNK);
 *
 * This invocation results in the following static functions:
 *
 *     static inline void initialize_my_fancy_type_allocator(void);
 *     static inline my_fancy_type *allocate_my_fancy_type(void);
 *     static inline void deallocate_my_fancy_type(my_fancy_type *object);
 *     static inline void deinitialize_my_fancy_type_allocator(void);
 *
 * Low-level operations can still be used if desired, via the provided:
 *
 *     static op_allocator my_fancy_type_allocator;
 *
 * Use of all of these is precisely the same as for the low-level interface, but
 * the allocation and deallocation is type-safe.
 *
 * @note There is no real need to call the initialization function because the
 *       allocation function does that automatically on first invocation.
 *
 * @note The deinitialization function is really only needed for obsessive
 *       clean-up.
 *
 * @{
 */

/** @brief Declare a component allocator.
 *
 */
#define OP_HL_DECLARE_ALLOCATOR(TYPE, COUNT, MODE)                            \
static op_allocator TYPE##_allocator;                                         \
static bool TYPE##_allocator_initialized = false;                             \
static inline void initialize_##TYPE##_allocator(void)                        \
{ TYPE##_allocator = op_ll_initialize_allocator(sizeof(TYPE), COUNT, MODE); } \
static inline void deinitialize_##TYPE##_allocator(void)                      \
{ op_ll_deinitialize_allocator(TYPE##_allocator); }                           \
static inline TYPE *allocate_##TYPE(void)                                     \
{ if (!TYPE##_allocator_initialized)                                          \
  { initialize_##TYPE##_allocator(); TYPE##_allocator_initialized = true; }   \
  return op_ll_allocate_object(TYPE##_allocator); }                           \
static inline void deallocate_##TYPE(TYPE *object)                            \
{ UNUSED(deinitialize_##TYPE##_allocator);                                    \
op_ll_deallocate_object(TYPE##_allocator, object); }

/**@}*/

#ifdef OPALLOC_IMPLEMENTATION
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

// macro to make output strings easier
#ifndef OP_OS
#include <stdio.h>
#define OP_OS(...) printf(__VA_ARGS__)
#endif

/*******************************************************************************
* Static helper function declarations
*******************************************************************************/

static bool fill_chunks(op_allocator allocator, const size_t offset, const size_t object_count);
static bool grow_pool(op_allocator allocator);

/*******************************************************************************
* Opaque data structures
*******************************************************************************/

typedef enum _active
{
    NOT_IN_USE,
    IN_USE,
} _active;

typedef struct _ab_t
{
    union
    {
        _active  in_use;
        uint32_t force_alignment;   // enums aren't always 32-bit
    };
    uint8_t data[];
} _ab_t;

struct _op_allocator
{
    size_t object_size;
    size_t initial_count;
    size_t maximum_objects;
    bool   use_chunks;
    bool   use_linear;
    bool   initialized;
    _ab_t   **pool;
};

/*******************************************************************************
* Low-level API function definitions
*******************************************************************************/

void *op_ll_allocate_object(op_allocator allocator)
{
    void *rv = NULL;

    if (allocator && allocator->initialized)
    {
retry:

        for (size_t i = 0; i < allocator->maximum_objects; i++)
        {
            if (allocator->pool[i] == NULL)
            {
                allocator->pool[i] = calloc(1, sizeof(_ab_t) + allocator->object_size);

                if (allocator->pool[i])
                {
                    rv = allocator->pool[i]->data;
                    memset(rv, 0, allocator->object_size);
                    allocator->pool[i]->in_use = IN_USE;
                    break;
                }
                else
                {
                    op_error_handler(__FILE__, __LINE__, "Could not allocate desired object.");
                }
            }
            else if (allocator->pool[i]->in_use == NOT_IN_USE)
            {
                rv = allocator->pool[i]->data;
                memset(rv, 0, allocator->object_size);
                allocator->pool[i]->in_use = IN_USE;
                break;
            }
        }

        if (rv == NULL)
        {
            if (grow_pool(allocator))
            {
                /* This is safe.  If grow_pool() returns true, we are guaranteed there is space, so it won't loop. */
                goto retry;
            }
            else
            {
                op_error_handler(__FILE__, __LINE__, "Unable to grow allocation pool.");
            }
        }
    }
    else
    {
        op_error_handler(__FILE__, __LINE__, "Attempted to allocate from uninitialized allocator.");
    }

    if (rv == NULL)
    {
        op_error_handler(__FILE__, __LINE__, "Could not allocate desired object.");
    }

    return rv;
}

void op_ll_deallocate_object(const op_allocator allocator, const void *object)
{
    if (allocator != NULL && object != NULL)
    {
        for (size_t i = 0; i < allocator->maximum_objects; i++)
        {
            if (allocator->pool[i]->data == object)
            {
                allocator->pool[i]->in_use = NOT_IN_USE;
                return;
            }
        }
    }
    else
    {
        op_error_handler(__FILE__, __LINE__, "Invalid allocator or object while deallocating.");
    }
}

void op_ll_deinitialize_allocator(op_allocator allocator)
{
    if (allocator && allocator->initialized)
    {
        allocator->initialized = false;

        if (allocator->use_chunks)
        {
            if (allocator->use_linear)
            {
                /* this allocator is chunked and uses groups of initial_count allocations */
                for (size_t i = 0; i < allocator->maximum_objects; i += allocator->initial_count)
                {
                    free(allocator->pool[i]);
                }
            }
            else
            {
                /* this allocator is chunked and uses the doubling allocation scheme */
                free(allocator->pool[0]);   /* delete the first chunk of initial_count */

                /* each subsequent chunk starts at doublings from initial_count */
                for (size_t i = allocator->initial_count; i < allocator->maximum_objects; i <<= 1)
                {
                    free(allocator->pool[i]);
                }
            }
        }
        else
        {
            for (size_t i = 0; i < allocator->maximum_objects; i++)
            {
                free(allocator->pool[i]);   /* it's OK to free a NULL, and we initialize to NULL */
            }
        }

        free(allocator->pool);
        free(allocator);
    }
    else
    {
        op_error_handler(__FILE__, __LINE__, "Attempted to deinitialize an allocator that was not initialized.");
    }
}

op_allocator_stats op_ll_get_allocator_stats(const op_allocator allocator)
{
    op_allocator_stats rv = { 0 };

    if (allocator != NULL && allocator->initialized)
    {
        rv.object_size = allocator->object_size;
        rv.maximum_objects = allocator->maximum_objects;
        rv.active_objects = 0;

        for (size_t i = 0; i < allocator->maximum_objects && allocator->pool[i] != NULL; i++)
        {
            if (allocator->pool[i]->in_use == IN_USE)
            {
                rv.active_objects++;
            }
        }
    }

    return rv;
}

op_allocator op_ll_initialize_allocator(const size_t object_size, const size_t initial_count,
                                        const op_ll_allocator_mode mode)
{
    bool use_chunks = false, use_linear = false;

    switch (mode)
    {
    case OP_DOUBLING_INDIVIDUAL:
        use_chunks = false;
        use_linear = false;
        break;

    case OP_DOUBLING_CHUNK:
        use_chunks = true;
        use_linear = false;
        break;

    case OP_LINEAR_INDIVIDUAL:
        use_chunks = false;
        use_linear = true;
        break;

    case OP_LINEAR_CHUNK:
        use_chunks = true;
        use_linear = true;
        break;
    };

    op_allocator rv = calloc(1, sizeof(struct _op_allocator));

    if (rv)
    {
        /* set the allocator configuration */
        rv->object_size = object_size;
        rv->initial_count = rv->maximum_objects = initial_count;
        rv->use_chunks = use_chunks;
        rv->use_linear = use_linear;
        rv->initialized = true;

        /* allocate the space for the object pool itself. */
        if ((rv->pool = calloc(rv->maximum_objects, sizeof(_ab_t *))))
        {
            if (use_chunks)
            {
                if (!fill_chunks(rv, 0, rv->maximum_objects))
                {
                    free(rv->pool);
                    free(rv);
                    rv = NULL;
                }
            }
        }
        else
        {
            free(rv);
            rv = NULL;
            op_error_handler(__FILE__, __LINE__, "Could not allocate internal space for allocator.");
        }
    }
    else
    {
        op_error_handler(__FILE__, __LINE__, "Could not allocate space for allocator handle.");
    }

    return rv;
}

/*******************************************************************************
* Static helper function definitions
*******************************************************************************/

static bool fill_chunks(op_allocator allocator, const size_t offset, const size_t object_count)
{
    /* fill_chunks status
     *  - understanding: optional
     *  - breaking if you fuck with it: inevitable
     *
     * the test suite is watching you and will be VERY disappointed
     */
    bool rv = true;

    size_t entry_size = sizeof(_ab_t) + allocator->object_size;
    size_t chunk_size = object_count * entry_size;
    uint8_t *chunk;
    chunk = calloc(object_count, chunk_size);

    if (chunk)
    {
        for (size_t i = 0; i < object_count; i++)
        {
            allocator->pool[i + offset] = (_ab_t *) &chunk[i * entry_size];
        }
    }
    else
    {
        op_error_handler(__FILE__, __LINE__, "Could not fill chunks when initializing or expanding.");
        rv = false;
    }

    return rv;
}

static bool grow_pool(op_allocator allocator)
{
    bool rv = false;

    size_t old_size = allocator->maximum_objects;
    size_t new_size, grow_size;

    if (allocator->use_linear)
    {
        grow_size = allocator->initial_count;
        new_size = old_size + grow_size;
    }
    else
    {
        grow_size = old_size;
        new_size = old_size * 2;
    }

    allocator->pool = realloc(allocator->pool, sizeof(_ab_t *) * new_size);

    if (allocator->pool)
    {
        /* we use the for loop here because valgrind is too stupid to figure out memset */
        for (size_t i = old_size; i < new_size; i++)
        {
            allocator->pool[i] = NULL;
        }

        allocator->maximum_objects = new_size;

        if (allocator->use_chunks)
        {
            rv = fill_chunks(allocator, old_size, grow_size);
        }
        else
        {
            rv = true;
        }
    }
    else
    {
        op_error_handler(__FILE__, __LINE__, "Could not realloc while growing pool.");
    }

    return rv;
}

void default_op_error_handler(const char *file, const int line, const char *error_message)
{
    UNUSED(file);
    UNUSED(line);
    UNUSED(error_message);
}
op_error_handler_t op_error_handler = default_op_error_handler;

#endif //OPALLOC_IMPLEMENTATION

#endif //!OPALLOC_INCLUDED
